import getpass
import os
from typing import List
from langchain_nomic import NomicEmbeddings
from langchain_chroma import Chroma


def list_png_images(directory: str) -> List[str]:
    """
    List all PNG image paths from the specified directory and its subdirectories.

    Args:
        directory (str): The directory to search for PNG images.

    Returns:
        List[str]: A list of paths to PNG images.
    """
    png_files = []
    for root, _, files in os.walk(directory):
        for file in files:
            if file.endswith('.png'):
                png_files.append(os.path.join(root, file).replace("\\", "/"))
    return png_files


def load_img_db(db: str):
    '''
    Load or new one

    db: str, path to the directory where the db is saved
    '''
    ImageEmbed = NomicEmbeddings(
        model="_", vision_model="nomic-embed-vision-v1.5")

    vector_store = Chroma(
        collection_name="dunhuang_db",
        embedding_function=ImageEmbed,
        persist_directory=db,
    )
    return vector_store


def append_img_db(db: str, img_dir: str):
    """
    db: str, path to the directory where the db is saved
    img_dir: str, path to the directory where the appended images are saved
    """
    vector_store = load_img_db(db)
    # load images
    png_images = list_png_images(img_dir)
    # print(png_images)

    # Index images
    metadata = [{'path': image} for image in png_images]
    vector_store.add_images(uris=png_images, metadatas=metadata)
    print(f"Append all images in {img_dir} to {vector_store._collection_name} done")
    return vector_store


def main():
    if not os.environ.get("NOMIC_API_KEY"):
        # os.environ["NOMIC_API_KEY"] = getpass.getpass("Enter API key for NOMIC: ")
        # todo: remove key
        os.environ["NOMIC_API_KEY"] = "nk-o09gqPOC9Qq-b7SCo5-tYNOvORBvBV-YUUomTg3fbjU"

    # load db or build it
    db_path = "./data/dunhuang_db"
    vector_store = None
    if os.path.exists(db_path) and os.path.isdir(db_path) and os.listdir(db_path):
        vector_store = load_img_db(db_path)
    else:
        img_dir = "./data/dunhuang_raw"
        vector_store = append_img_db(db_path, img_dir)

    # query image
    # query_image = "./data/dunhuang_raw/0023/莫高窟第023窟.png"
    query_image = "./data/dunhuang_raw/0254/莫高窟第254窟_主室_南壁.png"
    print("Query image:", query_image)
    res = vector_store.similarity_search_by_image_with_relevance_score(uri=query_image, k=3)
    for doc, score in res:
        print(f"Result img {doc.metadata}, relevance score {score}")


if __name__ == "__main__":
    main()
